﻿#include <iostream>
#include <vector>
#include <string>
#include <cstring>
#include <cstdint>
#include <cstdlib>
#include <ctime>
#include <sstream>
#include <fstream>
#include <cassert>
#include <map>
#if defined(_WIN32)
#   define WIN32_LEAN_AND_MEAN
#   include <windows.h>
#   include <winsock2.h>
#   include <ws2tcpip.h>
#elif defined(__unix__)
#   include <sys/types.h>
#   include <sys/socket.h>
#   include <unistd.h>
#   include <netdb.h>
#   include <arpa/inet.h>
#else
#   error "unsupport platform"
#endif

using namespace std;

struct DNS
{
    /**
     * @brief 事务 ID
     * DNS 报文的 ID 标识。对于请求报文和其对应的应答报文，该字段的值是相同的。
     * 通过它可以区分 DNS 应答报文是对哪个请求进行响应的。
     */
    uint16_t id;
    struct
    {
        /**
         * 返回码字段，表示响应的差错状态。
         *    0: 表示没有错误
         *    1: 表示报文格式错误（Format error），服务器不能理解请求的报文
         *    2: 表示域名服务器失败（Server failure），因为服务器的原因导致没办法处理这个请求
         *    3: 表示名字错误（Name Error），只有对授权域名解析服务器有意义，指出解析的域名不存在
         *    4: 表示查询类型不支持（Not Implemented），即域名服务器不支持查询类型
         *    5: 表示拒绝（Refused），一般是服务器由于设置的策略拒绝给出应答，如服务器不希望对某些请求者给出应答
         */
        uint16_t rcode : 4;
        uint16_t checkdisable : 1; // 0
        uint16_t authenticated : 1; // 0
        uint16_t z : 1; // 保留字段，在所有的请求和应答报文中，它的值必须为 0。
        uint16_t recavail : 1; // 可用递归。该字段只出现在响应报文中。当值为 1 时，表示服务器支持递归查询。
        uint16_t recdesired : 1; // 期望递归。该字段能在一个查询中设置，并在响应中返回。该标志告诉名称服务器必须处理这个查询，这种方式被称为一个递归查询。如果该位为 0，且被请求的名称服务器没有一个授权回答，它将返回一个能解答该查询的其他名称服务器列表。这种方式被称为迭代查询。
        uint16_t truncated : 1; // 表示是否被截断。值为 1 时，表示响应已超过 512 字节并已被截断，只返回前 512 个字节。
        uint16_t authoritative : 1; // 授权应答，该字段在响应报文中有效。值为 1 时，表示名称服务器是权威服务器；值为 0 时，表示不是权威服务器。
        uint16_t opcode : 4; // 操作码。其中，0 表示标准查询；1 表示反向查询；2 表示服务器状态请求。
        uint16_t response : 1; // 查询请求/响应的标志信息。查询请求时，值为 0；响应时，值为 1。
    } flags;
    uint16_t queries_count; // 问题计数：DNS 查询请求的数目
    uint16_t answers_count; // 回答资源记录数：DNS 响应的数目。
    uint16_t auth_rr_count; // 权威名称服务器计数：权威名称服务器的数目
    uint16_t add_rr_count; // 附加资源记录数：额外的记录数目（权威名称服务器对应 IP 地址的数目）
};

string get_str(const uint8_t *&ptr, const char *buf);

int main(int argc, char *argv[])
{
    if (argc != 3)
    {
        cerr << "Usage: "<< argv[0] << " <dns server ip> <to find host name>" << endl;
        return 1;
    }
#ifdef _WIN32
    WSADATA wsa;
    WSAStartup(MAKEWORD(2, 2), &wsa);
#endif
    auto sock = ::socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    sockaddr_in addr = {0};
    addr.sin_family = AF_INET;
    inet_pton(AF_INET, argv[1], &addr.sin_addr.s_addr);
    addr.sin_port = htons(53);
    ::connect(sock, reinterpret_cast<sockaddr*>(&addr), sizeof(addr));

    size_t query_size = sizeof(DNS) + strlen(argv[2]) + 6;

    string pbuf(query_size, '\0');
    DNS *dns_query = reinterpret_cast<DNS*>(&pbuf[0]);
    dns_query->id = htons(1);
    dns_query->queries_count = htons(1);
    uint8_t *ptr = reinterpret_cast<uint8_t*>(&pbuf[sizeof(DNS)]);

    string s;
    for (istringstream iss(argv[2]); getline(iss, s, '.');)
    {
        *ptr++ = (uint8_t)s.size();
        memcpy(ptr, s.c_str(), s.size());
        ptr += s.size();
    }
    *ptr++ = 0;
    uint16_t *ptr1 = reinterpret_cast<uint16_t *>(ptr);
    ptr1[0] = htons(1); // type: A (Host Address)
    ptr1[1] = htons(1); // Class: IN

    ::send(sock, pbuf.data(), (int)pbuf.size(), 0);
    char buf[0x1000];
    auto answer_size = ::recv(sock, buf, sizeof buf, 0);
#if defined(_WIN32)
    closesocket(sock);
#elif defined(__unix__)
    close(sock);
#else
#   error "unsupport platform"
#endif
    assert(answer_size > query_size);
    const DNS *dns_answer = reinterpret_cast<const DNS*>(buf);
    assert(dns_answer->id == dns_query->id);
    assert(dns_answer->queries_count == dns_query->queries_count);

    const uint8_t *ptr2 = reinterpret_cast<const uint8_t*>(buf + query_size);

    for (size_t i = 0; i < ntohs(dns_answer->answers_count); i++)
    {
        cout << i << endl;
        cout << "\tname: " << get_str(ptr2, buf) << endl;
        cout << "\ttype: " << ntohs(*reinterpret_cast<const uint16_t*>(ptr2)) << endl;
        ptr2 += 2;
        cout << "\tclass: " << ntohs(*reinterpret_cast<const uint16_t*>(ptr2)) << endl;
        ptr2 += 2;
        cout << "\ttime to live: " << ntohl(*reinterpret_cast<const uint32_t*>(ptr2)) << endl;
        ptr2 += 4;
        uint16_t len = ntohs(*reinterpret_cast<const uint16_t*>(ptr2));
        ptr2 += 2;
        assert(len == 4);
        char addr_buf[0x10];
        cout << "\taddress: " << inet_ntop(AF_INET, ptr2, addr_buf, sizeof addr_buf) << endl;
        ptr2 += 4;
    }

#ifdef _WIN32
    WSACleanup();
#endif
    return 0;
}

string get_str(const uint8_t *&ptr, const char *buf)
{
    ostringstream oss;
    while(*ptr)
    {
        if(!oss.str().empty())
            oss << '.';
        if((*ptr & 0xC0) == 0xC0)
        {
            uint16_t off = ntohs(*reinterpret_cast<const uint16_t*>(ptr)) & 0x3FFF;
            const uint8_t *p = reinterpret_cast<const uint8_t*>(buf + off);
            ptr += 2;
            oss << get_str(p, buf);
            return oss.str();
        }
        else
        {
            const char *p = reinterpret_cast<const char*>(ptr + 1);
            oss << string(p, p + *ptr);
            ptr += *ptr + 1;
        }
    }
    ptr++;
    return oss.str();
}
